import rospy
import std_msgs.msg
from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import numpy as np


import struct
import std_msgs
from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2

import pyrealsense2 as rs

intrinsics_matrix = np.array([[609.0, 0., 320.],
                         [ 0.,609.0, 240],
                         [0,0,1]], dtype=np.float32)

factor_depth = 1000

class CameraInfo():
    def __init__(self, width, height, fx, fy, cx, cy, scale):
        self.width = width
        self.height = height
        self.fx = fx
        self.fy = fy
        self.cx = cx
        self.cy = cy
        self.scale = scale


def convert_numpy_to_pointcloud2(np_array):
    """
    Convert numpy array to sensor_msgs/PointCloud2.
    
    np_array: A numpy array of shape (n, 6) with each row as (x, y, z, r, g, b)
    
    Returns:
    PointCloud2 message
    """
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = "map"

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgba', 12, PointField.UINT32, 1),
    ]

    point_cloud = []
    for point in np_array:
        x, y, z, r, g, b = point
        r = min(max(int(r), 0), 255)
        g = min(max(int(g), 0), 255)
        b = min(max(int(b), 0), 255)
        # a = 255  # Full opacity
        rgba = struct.unpack('I', struct.pack('BBBB', b, g, r, 255))[0]
        point_cloud.append([x, y, z, rgba])

    pc2_msg = pc2.create_cloud(header, fields, point_cloud)
    
    return pc2_msg

def rgbd2cloud(color, depth, camera, rgb_flag = False, num_point = 10000):
    assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width)
    xmap = np.arange(camera.width)
    ymap = np.arange(camera.height)
    xmap, ymap = np.meshgrid(xmap, ymap)  # 创建图像网格
    points_z = depth / camera.scale       # 深度信息  使用了广播机制
   
    points_x = (xmap - camera.cx) * points_z / camera.fx  # X = (u - cx) * Z /fx
    points_y = (ymap - camera.cy) * points_z / camera.fy  # Y = (v - cx) * Y /fy
    cloud = np.stack([points_x, points_y, points_z], axis=-1)  # 沿着最后一个轴增加维度
    if rgb_flag:
        color = color[:, :, ::-1]   # bgr2rgb
        cloud = np.concatenate((cloud, color), axis=2)
        # print("cloud: ", cloud.shape)
        cloud = cloud.reshape([-1, 6])
    else:
        cloud = cloud.reshape([-1, 3])
    cloud = cloud[(cloud[:, 2] > 0.2) & (cloud[:, 2] < 2)]
    
    cloud_size = cloud.shape[0]
    if(cloud_size >= num_point):
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    else:
        idxs1 = np.arange(cloud_size)  # 补点
        idxs2 = np.random.choice(cloud_size, num_point - cloud_size, replace=True)
        idxs = np.concatenate([idxs1, idxs2], axis=0)
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    cloud = cloud[idxs]
    
    return cloud


 # 初始化深度相机
pipeline = rs.pipeline()
config = rs.config()

# 配置深度和颜色流
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

# 开始更新配置和同步设置
pipeline.start(config)
align = rs.align(rs.stream.color)  # Create align object for depth-color alignment

camera = CameraInfo(640, 480, intrinsics_matrix[0][0], intrinsics_matrix[1][1], intrinsics_matrix[0][2], intrinsics_matrix[1][2], factor_depth)
        

rospy.init_node("grasp_node")
rate = rospy.Rate(10)
pub_cloud = rospy.Publisher("/points", PointCloud2, queue_size=10)

while not rospy.is_shutdown():
# 等待深度数据帧和RGB数据帧，设置等待时间为10秒
    frames = pipeline.wait_for_frames(timeout_ms=10000)
    aligned_frames = align.process(frames)
    if not aligned_frames:
        continue  # If alignment fails, go bac

    depth_frame = frames.get_depth_frame()
    color_frame = frames.get_color_frame()

    if not depth_frame or not color_frame:
        continue
    # 获取深度图像的原始数据
    depth_data = np.asanyarray(depth_frame.get_data())
    # 获取RGB图像的原始数据
    color_data = np.asanyarray(color_frame.get_data())

    cloud = rgbd2cloud(color_data, depth_data, camera, True, 10000)
    print(cloud.shape)
    msg_cloud = convert_numpy_to_pointcloud2(cloud)
    pub_cloud.publish(msg_cloud)
    rate.sleep()
